import os
import cv2 as cv
import numpy as np
import json
import shutil
import PIL
from PIL import Image, ImageDraw
from sklearn.cluster import KMeans
import glob
import scipy.misc
import imageio
import skimage.exposure
import skimage.feature
import skimage.filters
import skimage.io
import skimage.morphology
import sys


def find_seed(img, n_split=5, sampler_factor=3.0, thresh_pixel=10, n_vsplit=3):
    h, w = img.shape
    pos_l = w // 2
    pos_r = pos_l
    step = int(np.round(w / (n_split * sampler_factor)))

    def worker(img, sampler_pos):
        mask = np.zeros((h, w), dtype=np.uint8)
        mask[:, sampler_pos] = 255
        res = cv.bitwise_and(img, img, mask=mask)
        res_middle = res[:, sampler_pos]
        idx = res_middle != 0
        col_y = np.arange(0, h, dtype=int)
        dots_y = col_y[idx]
        dots_y = dots_y.reshape(-1, 1)
        kmeans = KMeans(n_clusters=n_vsplit).fit(dots_y)
        centers = kmeans.cluster_centers_
        centers = np.sort(centers, axis=0).ravel()
        if n_vsplit == 3:
            for diff in [centers[2] - centers[1], centers[1] - centers[0]]:
                if diff < thresh_pixel:
                    raise ValueError(f'diff={diff} < thresh={thresh_pixel}')
            point1 = (sampler_pos, int(np.round((centers[0] + centers[1]) / 2)))
            point2 = (sampler_pos, int(np.round((centers[2] + centers[1]) / 2)))
            return point1, point2
        elif n_vsplit == 2:
            for diff in [centers[1] - centers[0]]:
                if diff < thresh_pixel:
                    raise ValueError(f'diff={diff} < thresh={thresh_pixel}')
            point1 = (sampler_pos, int(np.round((centers[0] + centers[1]) / 2)))
            return point1, None
        else:
            raise Exception(f'n_vsplit={n_vsplit} is not supported currently!')

    pos_arr = []
    for i in range(w // step):
        if pos_l >= 0:
            pos_arr.append(pos_l)
        if pos_r <= w - 1:
            pos_arr.append(pos_r)
        pos_l -= step
        pos_r += step
    pos_arr = pos_arr[1:]
    # print(len(pos_arr), pos_arr)  # tmp
    for pos in pos_arr:
        # print(pos)  # tmp
        try:
            p1, p2 = worker(img, pos)
            return p1, p2
            break
        except ValueError as ve:
            pass
    return None, None


def flood_fill(img, point):
    h, w = img.shape
    img3c = np.stack((img, img, img), axis=2)
    mask = np.zeros((h + 2, w + 2), dtype=np.uint8)
    cv.floodFill(img3c, mask, point, (0, 255, 0), 1, 1)
    unit_x = w // 5
    mask_tpl = np.zeros_like(img, dtype=np.uint8)
    area_list = []
    for i in range(5):
        xleft = i * unit_x
        mask = mask_tpl.copy()
        mask[:, xleft:xleft + unit_x] = 255
        res = cv.bitwise_and(img3c, img3c, mask=mask)
        idx = (res == (255, 255, 255))
        res[idx[:, :, 0]] = (0, 0, 0)
        ret, res = cv.threshold(res[:, :, 1], 127, 255, cv.THRESH_BINARY)
        area = cv.countNonZero(res)
        area_list.append(area)
    return area_list


def valid_all_areas(all_areas, count_thresh=3):
    areas = sum(all_areas, [])
    n = len(areas)
    if n == 5:
        return True
    elif n != 10:
        return False
    else:
        arr = np.array(areas, dtype=int)
        bool_arr = arr <= 0
        xcount = bool_arr.astype(int).sum()
        if xcount >= count_thresh:
            return False
        else:
            return True


def print_result(output, i, all_areas):
    areas = sum(all_areas, [])
    # print(areas)  # tmp

    n = len(areas)
    if n == 10:
        s = "%s,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d" % (
            i,
            areas[0] + areas[5],
            areas[1] + areas[6],
            areas[2] + areas[7],
            areas[3] + areas[8],
            areas[4] + areas[9],
            areas[0],
            areas[1],
            areas[2],
            areas[3],
            areas[4],
            areas[5],
            areas[6],
            areas[7],
            areas[8],
            areas[9]
        )
    elif n == 5:
        s = "%s,%d,%d,%d,%d,%d" % (
            i,
            areas[0],
            areas[1],
            areas[2],
            areas[3],
            areas[4],
        )
    else:
        raise Exception(f'n={n} in print_result!')
    # print(s)  # tmp
    # return  # tmp
    output.write(s)
    output.write("\n")


for path in [
    # "/home/data7/Dataset/Ophthalmology_data/20220509-OK/OK镜",
    # "/home/data7/Dataset/Ophthalmology_data/20220510-OK/OK3月",
    # "/home/data7/Dataset/Ophthalmology_data/20220510-OK/OK6月",
    # "/home/data7/Dataset/Ophthalmology_data/20220510-OK/OK初查",
    # "/home/data7/Dataset/Ophthalmology_data/20220510-软镜/软镜3月",
    # "/home/data7/Dataset/Ophthalmology_data/20220510-软镜/软镜6月",
    # "/home/data7/Dataset/Ophthalmology_data/20220510-软镜/软镜初查",
    # "/home/data7/Dataset/Ophthalmology_data/20220511-NewMGD",
    # "/home/data7/Dataset/Ophthalmology_data/20220517-newnormal-fiveequalparts",
    # "/home/data7/Dataset/Ophthalmology_data/正常下睑"
    # "/home/data7/Dataset/Ophthalmology_data/不均匀萎缩分级（随机）/MGD组",
    # "/home/data7/Dataset/Ophthalmology_data/不均匀萎缩分级（随机）/正常组",
    "/home/data7/Dataset/Ophthalmology_data/不均匀萎缩分级（随机）/螨虫组"
]:
    if not os.path.exists(os.path.join(path, "temp")):
        os.mkdir(os.path.join(path, "temp"))
    if not os.path.exists(os.path.join(path, "chart")):
        os.mkdir(os.path.join(path, "chart"))

    with open("%s/区域.json" % path) as f:
        chart = json.load(f)

    output = open('%s/uas.csv' % path, 'w', encoding='gbk', newline="")
    output.write('文件名,a,b,c,d,e,a1,b1,c1,d1,e1,a2,b2,c2,d2,e2\n')

    for k, v in chart.items():
        fnames = k.split(".")[0]
        # print(fnames)

        try:
            im = Image.open("%s/原图/%s.BMP" % (path, fnames))
            mask = Image.new(im.mode, im.size)
            draw = ImageDraw.Draw(mask)
            for r in v['regions']:
                points = list(zip(r['shape_attributes']['all_points_x'], r['shape_attributes']['all_points_y']))
                draw.polygon(points, outline='white', fill=None)

            im_chart = Image.new(im.mode, im.size)
            draw2 = ImageDraw.Draw(im_chart)
            for r in chart[k]['regions']:
                points = list(zip(r['shape_attributes']['all_points_x'], r['shape_attributes']['all_points_y']))
                draw2.polygon(points, outline='white', fill=None)

            ones = np.nonzero(mask)
            top, bottom = np.min(ones[0]), np.max(ones[0])
            left, right = np.min(ones[1]), np.max(ones[1])

            im = np.array(im)
            im_crop = im[top: bottom, left: right]
            # skimage.io.imsave("%s/images/%s.png" % (path, fnames), im_crop)

            im_chart = np.array(im_chart)
            # print(im_chart)
            chart_crop = im_chart[top: bottom, left: right]
            skimage.io.imsave("%s/temp/%s.png" % (path, fnames), chart_crop)

        except Exception as e:
            print(e, k)
            continue

    files = glob.glob("%s/temp/*.png" % path)
    files = sorted(files)

    for fpath in files:
        basename = os.path.basename(fpath)
        name = basename.split(".")[0]
        img = cv.imread(fpath, cv.IMREAD_GRAYSCALE)
        ret, img = cv.threshold(img, 1, 255, cv.THRESH_BINARY)
        img_1 = img.copy()
        img_2 = img.copy()

        h, w = img.shape
        # cv.line(img, (0, h // 2), (w, h // 2), (255, 0, 0), 2)
        for i in range(5):
            i = i + 1
            num = w // 5
            cv.line(img, (num * i, 0), (num * i, h), (255, 0, 0), 2)

        cv.imwrite("%s/chart/%s" % (path, basename), img)


        points = []
        point1, point2 = find_seed(img_1)
        # print('pt1, pt2:', point1, point2)  # tmp
        is_only_double_lines = False
        if point1 is None:
            is_only_double_lines = True
            point1, point2 = find_seed(img_1, n_vsplit=2)
            # print('n_vsplit=2: pt1, pt2:', point1, point2)  # tmp
        if point1 is None and point2 is None:
            # ATTENTION: now you catch the error, maybe you should log it or report it.
            pass
        points.append(point1)
        points.append(point2)

        def get_areas(points):
            all_areas = []
            for p in points:
                if p is None:
                    continue
                areas = flood_fill(img_2, p)
                all_areas.append(areas)
            return all_areas
        all_areas = get_areas(points)
        # print(basename, all_areas)  # tmp
        if not is_only_double_lines:
            res = valid_all_areas(all_areas)
            if not res:
                # print(fname, 'RE-DO', file=sys.stderr)  # tmp
                point1, point2 = find_seed(img_1, n_vsplit=2)
                points = [point1, point2]
                all_areas = get_areas(points)
                # print(basename, all_areas)  # tmp


        print_result(output, name, all_areas)

    temp_path = os.path.join(path, "temp")
    shutil.rmtree(temp_path)







